import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression

data = np.loadtxt(r'../data/ex2data1.txt', delimiter=',')
x = data[:, :-1]
y = data[:, -1]

x = StandardScaler().fit_transform(x)

x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=666)

model = LogisticRegression(solver='liblinear')
model.fit(x_train, y_train)
s_train = model.score(x_train, y_train)
print(f'Training score = {s_train}')
s_test = model.score(x_test, y_test)
print(f'Testing score = {s_test}')

plt.scatter(x[:, 0], x[:, 1], c=y, s=1, zorder=100, cmap=plt.cm.get_cmap('rainbow'))

x_plt = np.array([x[:, 0].min(), x[:, 0].max()])


def get_y_plt(x_plt, model):
    return - (model.intercept_ + model.coef_[0][0] * x_plt) / model.coef_[0][1]


y_plt = get_y_plt(x_plt, model)
plt.plot(x_plt, y_plt, 'g-', label='border')
plt.legend()
plt.grid()
plt.show()